import asyncio
import logging
import socket
from command import CMD

from config import config
from eth_reply import *
from worker import HANDLERS, WORKER


logger = logging.getLogger()

class udp_protocol(asyncio.DatagramProtocol):
    def __init__(self, cmd_queue:asyncio.Queue, on_con_lost:asyncio.Future) -> None:
        self.on_con_lost = on_con_lost
        self._cmd_queue = cmd_queue
    
    def connection_made(self, transport):
        self.transport = transport
        logger.info('UDP Connected.')


    def datagram_received(self, data, addr):
        try:
            self._cmd_queue.put_nowait({"addr":addr, "data":data})
        except asyncio.QueueFull:
            logger.critical("RX_Q is Full")

    def error_received(self, exc):
        logger.error('UDP Error received:', exc)
        try:
            self.on_con_lost.set_result(True)
        except asyncio.InvalidStateError as e:
            logger.debug(e, exc_info=True)

    def connection_lost(self, exc):
        logger.error("UDP Connection closed")
        try:
            self.on_con_lost.set_result(True)
        except asyncio.InvalidStateError as e:
            logger.debug(e, exc_info=True)

    def close(self):
        self.transport.close()
    
    def __del__(self):
        self.close()


class UDP_CMD_WORKER(WORKER):

    handlers = HANDLERS()
    
    def __init__(self, cfg:config, socket:socket.socket) -> None:
        super().__init__()
        self.cfg = cfg

        self.transport = None
        self.protocol : udp_protocol = None

        self.__s = socket

        self._loop = asyncio.get_event_loop()
        self._cmd_queue = asyncio.Queue(1024)

        self.tasks = [self.__connect_udp()]


    async def __connect_udp(self):
        loop = self._loop
        while True:
            try:
                logger.info('Connecting UDP....')

                self.transport, self.protocol = await loop.create_datagram_endpoint(
                    lambda: udp_protocol(self._cmd_queue, on_con_lost=loop.create_future()), 
                    sock=self.__s)

                send_recv_future = asyncio.ensure_future(self.__recv_handle_cmd())
                self.is_running = True

                await self.protocol.on_con_lost
                logger.error('UDP DISCONNECTED!')
                self.is_running = False
                if not send_recv_future.done():
                    send_recv_future.cancel()
                del self.protocol
            
            except OSError as e:
                logger.warning('Setted host IP may NOT right.')

            finally:
                await asyncio.sleep(5)


    async def __recv_handle_cmd(self):
        while True:
            try:
                rdata = await self._cmd_queue.get()
                data = rdata["data"]
                addr = rdata["addr"]
                logger.debug("Receive %i bytes form %r. Data : %s" % (len(data), addr, data))

                cmd = CMD.parse(data)

                reply_addr = (self.cfg.pc_ip, self.cfg.pc_port)

                if cmd:
                    logger.info("Parsed command:%s" % cmd)
                    try:
                        res = await cmd.handle()
                    except ValueError as e:
                        logger.error(e, exc_info=True)
                        res = [False]
                    # Only reply last handlers result.
                    reply = cmd.pack_reply(res[-1])
                    if reply:
                        logger.info("Reply %r for command:%s ====> %s" % (reply_addr, cmd, reply if len(reply) < 300 else (reply[:300] + b'...')))
                        self.__s.sendto(reply, reply_addr)

            except Exception as e:
                logger.exception(e)


    # Just for test.
    # @handlers(CMD.Type.DO_NOTHING)
    # def do_nothing(self, cmd):
    #     return True
        